package com.sun.btrace.client;

import java.io.Console;
import java.io.File;
import java.io.IOException;
import java.io.PrintWriter;
import java.net.Socket;
import java.net.URI;
import java.net.UnknownHostException;

import com.sun.btrace.client.Client;
import com.sun.tools.attach.VirtualMachine;

public class RegistAgent {
	public static volatile boolean exiting;
	public static final boolean DEBUG;
	public static final boolean TRACK_RETRANSFORM;
	public static final boolean UNSAFE;
	public static final boolean DUMP_CLASSES;
	public static final String DUMP_DIR;
	public static final String PROBE_DESC_PATH;
	public static final int BTRACE_DEFAULT_PORT = 2020;

	private static final Console con;
	private static final PrintWriter out;

	static {
		DEBUG = Boolean.getBoolean("com.sun.btrace.debug");
		if (isDebug())
			debugPrint("btrace debug mode is set");
		TRACK_RETRANSFORM = Boolean
				.getBoolean("com.sun.btrace.trackRetransforms");
		if (isDebug() && TRACK_RETRANSFORM)
			debugPrint("trackRetransforms flag is set");
		UNSAFE = Boolean.getBoolean("com.sun.btrace.unsafe");
		if (isDebug() && UNSAFE)
			debugPrint("btrace unsafe mode is set");
		DUMP_CLASSES = Boolean.getBoolean("com.sun.btrace.dumpClasses");
		if (isDebug() && DUMP_CLASSES)
			debugPrint("dumpClasses flag is set");
		DUMP_DIR = System.getProperty("com.sun.btrace.dumpDir", ".");
		if (DUMP_CLASSES) {
			if (isDebug())
				debugPrint("dumpDir is " + DUMP_DIR);
		}
		PROBE_DESC_PATH = System.getProperty("com.sun.btrace.probeDescPath",
				".");
		con = System.console();
		out = (con != null) ? con.writer() : new PrintWriter(System.out);
	}

	private static boolean isDebug() {
		return DEBUG;
	}

	private static void debugPrint(String msg) {
		System.out.println("DEBUG: " + msg);
	}

	public void attach(String pid) throws IOException {
		try {
			String agentPath = "/btrace-agent.jar";
			String tmp = Client.class.getClassLoader()
					.getResource("com/sun/btrace").toString();
			tmp = tmp.substring(0, tmp.indexOf("!"));
			tmp = tmp.substring("jar:".length(), tmp.lastIndexOf("/"));
			agentPath = tmp + agentPath;
			agentPath = new File(new URI(agentPath)).getAbsolutePath();
			attach(pid, agentPath, null, null);
		} catch (RuntimeException re) {
			throw re;
		} catch (IOException ioexp) {
			throw ioexp;
		} catch (Exception exp) {
			throw new IOException(exp.getMessage());
		}
	}

	/**
	 * Attach the BTrace client to the given Java process. Loads BTrace agent on
	 * the target process if not loaded already. Accepts the full path of the
	 * btrace agent jar. Also, accepts system classpath and boot classpath
	 * optionally.
	 */
	public void attach(String pid, String agentPath, String sysCp, String bootCp)
			throws IOException {
		try {
			VirtualMachine vm = null;
			if (DEBUG) {
				debugPrint("attaching to " + pid);
			}
			vm = VirtualMachine.attach(pid);
			if ((DEBUG)) {
				debugPrint("checking port availability: " + BTRACE_DEFAULT_PORT);
			}
			int serverPort = Integer.parseInt(vm.getSystemProperties()
					.getProperty("btrace.port", "-1"));
			if (serverPort != -1) {
				if (serverPort != BTRACE_DEFAULT_PORT) {
					throw new IOException(
							"Can not attach to PID "
									+ pid
									+ " on port "
									+ BTRACE_DEFAULT_PORT
									+ ". There is already a BTrace server active on port "
									+ serverPort + "!");
				}
			} else {
				if (!isPortAvailable(BTRACE_DEFAULT_PORT)) {
					throw new IOException("Port " + BTRACE_DEFAULT_PORT
							+ " unavailable.");
				}
			}

			if (DEBUG) {
				debugPrint("attached to " + pid);
			}

			if (DEBUG) {
				debugPrint("loading " + agentPath);
			}
			String agentArgs = "port=" + BTRACE_DEFAULT_PORT;
			if (DEBUG) {
				agentArgs += ",debug=true";
			}
			if (UNSAFE) {
				agentArgs += ",unsafe=true";
			}
			if (DUMP_CLASSES) {
				agentArgs += ",dumpClasses=true";
				agentArgs += ",dumpDir=" + DUMP_DIR;
			}
			if (TRACK_RETRANSFORM) {
				agentArgs += ",trackRetransforms=true";
			}
			if (bootCp != null) {
				agentArgs += ",bootClassPath=" + bootCp;
			}
			if (sysCp == null) {
				sysCp = getToolsJarPath();
			}
			agentArgs += ",systemClassPath=" + sysCp;
			agentArgs += ",probeDescPath=" + PROBE_DESC_PATH;
			if (DEBUG) {
				debugPrint("agent args: " + agentArgs);
			}
			vm.loadAgent(agentPath, agentArgs);
			if (DEBUG) {
				debugPrint("loaded " + agentPath);
			}
		} catch (RuntimeException re) {
			throw re;
		} catch (IOException ioexp) {
			throw ioexp;
		} catch (Exception exp) {
			throw new IOException(exp.getMessage());
		}
	}

	private String getToolsJarPath() {
		// try to get absolute path of tools.jar
		// first check this application's classpath
		String[] components = System.getProperty("java.class.path").split(
				File.pathSeparator);
		for (String c : components) {
			if (c.endsWith("tools.jar")) {
				return new File(c).getAbsolutePath();
			} else if (c.endsWith("classes.jar")) { // MacOS specific
				return new File(c).getAbsolutePath();
			}
		}
		// we didn't find -- make a guess! If this app is running on a JDK
		// rather
		// than a JRE there will be a tools.jar in $JDK_HOME/lib directory.
		if (System.getProperty("os.name").startsWith("Mac")) {
			String java_home = System.getProperty("java.home");
			String java_mac_home = java_home.substring(0,
					java_home.indexOf("/Home"));
			return java_mac_home + "/Classes/classes.jar";
		} else {
			return System.getProperty("java.home") + "../lib/tools.jar";
		}
	}

	private static boolean isPortAvailable(int port) {
		Socket clSocket = null;
		try {
			clSocket = new Socket("127.0.0.1", port);
		} catch (UnknownHostException ex) {
		} catch (IOException ex) {
			clSocket = null;
		}
		if (clSocket != null) {
			try {
				clSocket.close();
			} catch (IOException e) {
			}
			return false;
		}
		return true;
	}

	public static String help() {
		String h = "" + "should give pid to attach!!!";
		return h;
	}

	/**
	 * @param args
	 * @throws IOException
	 */
	public static void main(String[] args) throws IOException {
		if (args.length <= 0)
			System.out.println(RegistAgent.help());

		RegistAgent r = new RegistAgent();

		int port = -1;
		try {
			port = Integer.parseInt(args[0]);
		} catch (Exception e) {
			System.out.println(RegistAgent.help());
		}
		if (port <= 0)
			System.out.println(RegistAgent.help());

		r.attach(args[0]);

	}

}
